package main

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"sort"
	"time"
)

var (
	cmdOut bool
	cmdLog *log.Logger
)

func init() {
	f, err := os.OpenFile("commands.out", os.O_CREATE|os.O_APPEND|os.O_WRONLY, os.ModePerm)
	if err != nil {
		log.Fatalln(err)
	}
	cmdLog = log.New(io.MultiWriter(os.Stdout, f), "", 0)
}

func SetVirshDebug(b bool) {
	cmdOut = b
}

//获取正在运行中的所有虚拟机
func GetRunningVMs() (VMs []string, err error) {
	cmd := `virsh list | tail -n +3 | sed '/^$/d'` //从第3行开始输出，并且过滤空白行
	c := exec.Command("/bin/sh", "-c", cmd)
	out, err := c.CombinedOutput()
	if cmdOut {
		cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
		cmdLog.Println(string(out))
	}

	if err != nil {
		return nil, err
	}

	for _, s := range bytes.Split(out, []byte("\n")) {
		if bytes.TrimSpace(s) == nil {
			continue
		}
		fields := bytes.Fields(s)
		VMs = append(VMs, string(fields[1]))
	}
	return
}

//关闭运行中的虚拟机
func ShutdownVMs(VMs []string, force bool) {
	//优雅的关闭虚拟机
	for _, vm := range VMs {
		cmd := fmt.Sprintf("virsh shutdown --domain %q", vm)
		go func() {
			c := exec.Command("/bin/sh", "-c", cmd)
			out, _ := c.CombinedOutput()
			if cmdOut {
				cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
				cmdLog.Println(string(out))
			}
		}()
		time.Sleep(time.Second * 10)
	}
	time.Sleep(time.Minute * 10)

	//检查虚拟机是否已经关闭，如果没有则强制关闭
	if force {
		for _, vm := range VMs {
			cmd := fmt.Sprintf("virsh destroy --domain %q", vm)
			go func() {
				if !IsShutdown(vm) { //强制关闭
					c := exec.Command("/bin/sh", "-c", cmd)
					out, _ := c.CombinedOutput()
					if cmdOut {
						cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
						cmdLog.Println(string(out))
					}
				}
			}()
			time.Sleep(time.Second * 10)
		}
		time.Sleep(time.Minute * 5)
	}
}

//启动虚拟机
func StartVMs(VMs []string) {
	for _, vm := range VMs {
		cmd := fmt.Sprintf("virsh start --domain %q", vm)
		go func() {
			c := exec.Command("/bin/sh", "-c", cmd)
			out, _ := c.CombinedOutput()
			if cmdOut {
				cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
				cmdLog.Println(string(out))
			}
		}()
		time.Sleep(time.Second * 10)
	}
}

//检查虚拟机状态是否已经关闭
func IsShutdown(vm string) bool {
	cmd := fmt.Sprintf("virsh domstate --domain %q", vm)
	c := exec.Command("/bin/sh", "-c", cmd)
	out, err := c.CombinedOutput()
	if cmdOut {
		cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
		cmdLog.Println(string(out))
	}

	if err != nil || bytes.Contains(out, []byte("running")) {
		return false
	}

	return true
}

//获取虚拟机的快照
func GetSnapshots(vm, prefix string) (snaps []string, err error) {
	//从第3行开始输出，并且过滤空白行，进行日期和时间排序
	cmd := fmt.Sprintf("virsh snapshot-list --domain %q | tail -n +3 | sed '/^$/d'", vm)
	c := exec.Command("/bin/sh", "-c", cmd)
	out, err := c.CombinedOutput()
	if cmdOut {
		cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
		cmdLog.Println(string(out))
	}

	if err != nil {
		return nil, err
	}

	var ss = make([]snapshot, 0, 100)
	for _, s := range bytes.Split(out, []byte("\n")) {
		if bytes.TrimSpace(s) == nil {
			continue
		}
		fields := bytes.Fields(s)

		if bytes.HasPrefix(fields[0], []byte(prefix)) {
			//获取快照名称
			s := snapshot{Name: string(fields[0])}
			//获取快照时间
			s.Time, err = time.Parse("2006-01-02 15:04:05", fmt.Sprintf("%s %s", fields[1], fields[2]))
			if err != nil {
				return nil, err
			}
			ss = append(ss, s)
		}
	}

	//排序
	sort.Sort(snapshots(ss))
	for _, i := range ss {
		snaps = append(snaps, i.Name)
	}
	return
}

//删除一个快照
func DeleteSnapshot(vm, snapshotname string) error {
	cmd := fmt.Sprintf("virsh snapshot-delete --domain %q --snapshotname %q", vm, snapshotname)
	c := exec.Command("/bin/sh", "-c", cmd)
	out, err := c.CombinedOutput()
	if cmdOut {
		cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
		cmdLog.Println(string(out))
	}

	if err != nil {
		return err
	}

	if bytes.Contains(out, []byte("error")) { //标准输出提示的错误
		return errors.New(string(out))
	}

	return nil
}

//创建一个新的快照
func CreateSnapshot(vm, snapshotname string) error {
	cmd := fmt.Sprintf("virsh snapshot-create-as --domain %q --name %q", vm, snapshotname)
	c := exec.Command("/bin/sh", "-c", cmd)
	out, err := c.CombinedOutput()
	if cmdOut {
		cmdLog.Printf("Execut the %q at %s\n", cmd, time.Now().Format("2006/01/02 15:04:05"))
		cmdLog.Println(string(out))
	}

	if err != nil {
		return err
	}

	if bytes.Contains(out, []byte("error")) { //标准输出提示的错误
		return errors.New(string(out))
	}

	return nil
}

//快照排序
type snapshot struct {
	Name string
	Time time.Time
}

type snapshots []snapshot

func (s snapshots) Len() int           { return len(s) }
func (s snapshots) Less(i, j int) bool { return s[i].Time.Before(s[j].Time) }
func (s snapshots) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
